from paddleocr import PaddleOCR
from paddleocr import PPStructureV3
from fastapi import FastAPI, Body
from fastapi.responses import FileResponse
from fastapi.responses import StreamingResponse
import zipfile
from io import BytesIO
import requests
import os
import shutil
from enum import Enum

imgIndex = 0

app = FastAPI()

pipeline_OCR = PaddleOCR(
    use_doc_orientation_classify=False, 
    use_doc_unwarping=False, 
    use_textline_orientation=False)

pipeline_structure = PPStructureV3()

def prepare_output_dir(dir_path):
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
    os.makedirs(dir_path)


class DirType(str, Enum):
    ocr = "ocr"
    structure = "structure"

@app.get("/clear_dir")
def clear_dir():
    global imgIndex
    imgIndex = 0
    output_dir = "output_OCR"
    prepare_output_dir(output_dir)
    output_dir = "output_structure"
    prepare_output_dir(output_dir)
    return {
        "result": "清空目录"
    }

@app.post("/predict_OCR")
def predict_OCR(url: str):
    output_dir = "output_OCR"
    global imgIndex
    # 下载图片到本地
    local_path = f"img_{imgIndex}.png"
    url = f"http://10.100.3.76{url}"
    r = requests.get(url)
    with open(local_path, "wb") as f:
        f.write(r.content)
    result = pipeline_OCR.predict(local_path)
    for res in result:
        res.print()
        res.save_to_img(output_dir)
        res.save_to_json(output_dir)
    os.remove(local_path)
    imgIndex += 1
    return {
        "result": "OCR识别完了"
    }

@app.post("/predict_structure")
def predict_structure(url: str):
    output_dir = "output_structure"
    global imgIndex
    # 下载图片到本地
    local_path = f"img_{imgIndex}.png"
    url = f"http://10.100.3.76{url}"
    r = requests.get(url)
    with open(local_path, "wb") as f:
        f.write(r.content)
    result = pipeline_structure.predict(local_path)
    for res in result:
        res.print()
        res.save_to_json(output_dir)
        res.save_to_markdown(output_dir)
    os.remove(local_path)
    imgIndex += 1
    return {
        "result": "结构化识别完了"
    }

@app.get("/result")
def get_result():
    results = {
        "ocr_results": {},
        "structure_results": {}
    }
    ocr_dir = "output_OCR"
    structure_dir = "output_structure"

    if os.path.exists(ocr_dir):
        ocr_files = os.listdir(ocr_dir)
        results["ocr_results"]["png_files"] = [f for f in ocr_files if f.endswith(".png")]
        results["ocr_results"]["json_files"] = [f for f in ocr_files if f.endswith(".json")]

    if os.path.exists(structure_dir):
        structure_files = os.listdir(structure_dir)
        results["structure_results"]["md_files"] = [f for f in structure_files if f.endswith(".md")]
        results["structure_results"]["json_files"] = [f for f in structure_files if f.endswith(".json")]

    return results

@app.get("/output/{dir_type}/{filename}")
def get_output_file(dir_type: DirType, filename: str):
    if dir_type == "ocr":
        base_path = "output_OCR"
    elif dir_type == "structure":
        base_path = "output_structure"
    else:
        return {"error": "Invalid directory type"}

    file_path = os.path.join(base_path, filename)

    if not os.path.exists(file_path):
        return {"error": "File not found"}

    media_type = "application/octet-stream"
    if filename.endswith(".png"):
        media_type = "image/png"
    elif filename.endswith(".json"):
        media_type = "application/json"
    elif filename.endswith(".md"):
        media_type = "text/markdown"

    return FileResponse(file_path, media_type=media_type)

@app.get("/output")
def download_all_outputs():
    zip_buffer = BytesIO()
    with zipfile.ZipFile(zip_buffer, "w") as zipf:
        # 添加 output_OCR 文件夹下所有文件
        ocr_dir = "output_OCR"
        if os.path.exists(ocr_dir):
            for file in os.listdir(ocr_dir):
                file_path = os.path.join(ocr_dir, file)
                zipf.write(file_path, arcname=os.path.join("output_OCR", file))
        # 添加 output_structure 文件夹下所有文件
        structure_dir = "output_structure"
        if os.path.exists(structure_dir):
            for file in os.listdir(structure_dir):
                file_path = os.path.join(structure_dir, file)
                zipf.write(file_path, arcname=os.path.join("output_structure", file))
    zip_buffer.seek(0)
    return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=all_outputs.zip"})
